import math
import tkinter
# import matplotlib
# matplotlib.use('TkAgg')
import os
import warnings

import CR_DSPPytorch

os.environ["OMP_NUM_THREADS"] = "8"
from inspect import currentframe, getframeinfo
import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn as nn
import scipy.io as scio
from CR_DSPPytorch import EDCLayer, PhaseRecLayer, FIRLayer, format_rt, cmul, PerturbativeBlockLayer
import util
from TDCE import modified_filter_taps_by_cluster_real, modified_filter_taps_by_cluster_complex, quantization_filter_taps

os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
BASE_DIR = os.path.dirname(__file__)


def ERP_NLC_with_TDC(lp_range=[1], tap_list=[2048], init_method='LS_CO', save_result_flag=False, clusters_flag=False,
                     quantization_flag=False):
    frame = currentframe()
    experiment_name = getframeinfo(frame).function
    pick_syms_num = 131072
    modOrder = 4
    Qcache = np.zeros([len(lp_range), len(tap_list)])
    BERcache = np.zeros_like(Qcache)
    result_save_dir = os.path.join(BASE_DIR, f'result/{experiment_name}')
    if not os.path.exists(result_save_dir):
        os.makedirs(result_save_dir)
    chosen_device = 'cuda' if torch.cuda.is_available() else 'cpu'
    init_method = init_method
    save_result_flag = save_result_flag
    quantization_flag = quantization_flag

    for lpIndx, lp in enumerate(lp_range):
        for tapIndx, tap in enumerate(tap_list):
            print(f'lp={lp}, tap={tap}')
            # load data
            dataPath = os.path.join(BASE_DIR, 'data/simulation/16QAM32Gbaud3200kmHe/PMD_eliminated_by_CMA'
                                              f'/trSet_lp_{lp}.mat')
            data = scio.loadmat(dataPath)
            symbolRate = data['symbolRate'][0, 0]
            spanNum = data['spanNum'][0, 0]
            spanLen = data['spanLen'][0, 0]
            sig = data['sig']
            prbs = data['prbs']
            L = spanNum * spanLen
            D = 17e-6
            DL = D * L / 2

            sig = torch.from_numpy(sig[np.newaxis, ...])
            constellations = torch.from_numpy(util.CONST_16QAM)

            '''均衡色散所需总抽头书上界'''
            total_taps_limit = util.cal_tap_num_odd(DL=DL,
                                                    sample_freq=2 * symbolRate)
            print(f'taps upper bound: {total_taps_limit}')

            '''DSP procedure'''
            edc = EDCLayer(symbol_rate=symbolRate, DL=DL, case_num=2, tap=tap, sample_factor=2,
                           power_norm=True, init_method=init_method, time_domain=True)

            if clusters_flag:
                # edc.h.data = modified_filter_taps_by_cluster_in_complex_plane(edc.h.data, clusters=14)
                edc.h.data = modified_filter_taps_by_cluster_real(edc.h.data, clusters=14)
            if quantization_flag:
                edc.h.data = quantization_filter_taps(edc.h.data, bit_width=4)
            pr = PhaseRecLayer(1024)
            lms = FIRLayer(tap=32, case_num=2, power_norm=True, centor_one=True)

            sig = sig.to(chosen_device)
            lms = lms.to(chosen_device)
            edc = edc.to(chosen_device)

            sig = sig[..., 0:pick_syms_num * 2]
            prbs = prbs[..., 0:pick_syms_num * modOrder]

            sig = edc(torch.view_as_real(sig))
            sig = pr(sig[..., 1::2, :])
            lms.fit(sig, err_mode='DDM', constellations=constellations, iter_num=4, block_size=4028,
                    remain=2048, lr=5e-4)
            sig = lms(sig)

            sig = torch.view_as_complex(sig)
            sig = sig.cpu().data.numpy().squeeze()
            ber, Q = simulation_ERP_NLC_module(sig, prbs, lp, symbolRate, spanLen, spanNum)
            # sig, ber, _ = util.pr_ber(sig, prbs, constellations.cpu().data.numpy())

            # ber = np.mean(ber)
            BERcache[lpIndx, tapIndx] = ber
            Qcache[lpIndx, tapIndx] = Q
            print(f'lp: {lp}, tap:{tap}, BER: {ber}, Q: {Q}')

    if save_result_flag:
        with open(os.path.join(result_save_dir, f'ERP_NLC_TDC_vary_taps_12_17.csv'), 'w') as f:
            f.write('lp, tap, BER, Q\n')
            for lpIndx, lp in enumerate(lp_range):
                for tapIndx, tap in enumerate(tap_list):
                    f.write(f'{lp}, {tap}, {BERcache[lpIndx, tapIndx]}, {Qcache[lpIndx, tapIndx]}\n')
        print(f'result saved at \"{result_save_dir}\"')



def TDC_Cal_Complexity(lp=2, tap_list=[2048], cluster_list=[16], init_method='LS_CO', save_result_flag=False):
    frame = currentframe()
    experiment_name = getframeinfo(frame).function
    pick_syms_num = 131072
    modOrder = 4
    Qcache = np.zeros([len(tap_list), len(cluster_list)])
    BERcache = np.zeros_like(Qcache)
    Complexicache = np.zeros_like(Qcache)
    result_save_dir = os.path.join(BASE_DIR, f'result/ERP_NLC_with_TDC')
    if not os.path.exists(result_save_dir):
        os.makedirs(result_save_dir)
    chosen_device = 'cuda' if torch.cuda.is_available() else 'cpu'
    init_method = init_method
    save_result_flag = save_result_flag

    for tapIndx, tap in enumerate(tap_list):
        for clusterIndx, cluster in enumerate(cluster_list):
            # load data
            dataPath = os.path.join(BASE_DIR, 'data/simulation/16QAM32Gbaud3200kmHe/PMD_eliminated_by_CMA'
                                              f'/trSet_lp_{lp}.mat')
            data = scio.loadmat(dataPath)
            symbolRate = data['symbolRate'][0, 0]
            spanNum = data['spanNum'][0, 0]
            spanLen = data['spanLen'][0, 0]
            sig = data['sig']
            prbs = data['prbs']
            L = spanNum * spanLen
            D = 17e-6
            DL = D * L / 2

            sig = torch.from_numpy(sig[np.newaxis, ...])
            constellations = torch.from_numpy(util.CONST_16QAM)

            '''均衡色散所需总抽头书上界'''
            total_taps_limit = util.cal_tap_num_odd(DL=DL,
                                                    sample_freq=2 * symbolRate)
            print(f'taps upper bound: {total_taps_limit}')

            '''DSP procedure'''
            edc = EDCLayer(symbol_rate=symbolRate, DL=DL, case_num=2, tap=tap, sample_factor=2,
                           power_norm=True, init_method=init_method, time_domain=True)
            edc.h.data = modified_filter_taps_by_cluster_complex(edc.h.data, clusters=cluster)
            pr = PhaseRecLayer(1024)
            lms = FIRLayer(tap=32, case_num=2, power_norm=True, centor_one=True)

            sig = sig.to(chosen_device)
            lms = lms.to(chosen_device)
            edc = edc.to(chosen_device)

            sig = sig[..., 0:pick_syms_num * 2]
            prbs = prbs[..., 0:pick_syms_num * modOrder]

            sig = edc(torch.view_as_real(sig))
            sig = pr(sig[..., 1::2, :])
            lms.fit(sig, err_mode='DDM', constellations=constellations, iter_num=4, block_size=4028,
                    remain=2048, lr=5e-4)
            sig = lms(sig)

            sig = torch.view_as_complex(sig)
            sig = sig.cpu().data.numpy().squeeze()
            ber, Q = simulation_ERP_NLC_module(sig, prbs, lp, symbolRate, spanLen, spanNum)

            BERcache[tapIndx, clusterIndx] = ber
            Qcache[tapIndx, clusterIndx] = Q
            Complexicache[tapIndx, clusterIndx] = 4 * cluster
            print(f'lp: {lp}, tap:{tap}, BER: {ber}, Q: {util.ber2q(ber)}')

    if save_result_flag:
        with open(os.path.join(result_save_dir, f'Learned_TD_NLC_vary_tap_and_cluster.csv'), 'w') as f:
            f.write('tap,cluster,BER,Q,complexity\n')
            for tapIndx, tap in enumerate(tap_list):
                for clusterIndx, cluster in enumerate(cluster_list):
                    f.write(f' {tap}, {cluster}, {BERcache[tapIndx, clusterIndx]}, {Qcache[tapIndx, clusterIndx]}, '
                            f'{Complexicache[tapIndx, clusterIndx]}\n')
        print(f'result saved at \"{result_save_dir}\"')


def simulation_ERP_NLC_module(sig, prbs, lp, symbol_rate, span_length, span_number):
    '''
    sig: [2, 131072[, np.ndarray
    prbs: [2, 131072 * 4], np.ndarray
    '''
    from bayes_opt import BayesianOptimization
    from bayes_opt.util import UtilityFunction
    maxmn = 92
    symbol_num = sig.shape[-1]
    mod_order = 4

    cal_kwargs = {
        'Fs': symbol_rate,
        'sf': 20,
        'span_length': span_length,
        'span_number': span_number,
        'pre_cd': 0.5,
        'rolloff': 0.1,
        'maxmn': maxmn,
        'int_constraint': maxmn,
    }

    '''calculate effective length for ERP'''
    alp = 0.2e-3
    gm = 1.3e-3
    power = 10 ** (lp / 10 - 3)
    span_num = 40
    span_len = 80e3
    alpha = -np.log(np.power(10, -alp / 10))
    Leff = (1 - np.exp(- alpha * span_len * span_num)) / alpha

    perturbative_matrix_dir = os.path.join(BASE_DIR, f'perturbative_matrixes_sim')
    if not os.path.exists(perturbative_matrix_dir):
        os.makedirs(perturbative_matrix_dir)

    perturbative_matrix_path = os.path.join(perturbative_matrix_dir, f'win_size_{2 * maxmn + 1}.mat')

    prtMtrx = PerturbativeBlockLayer.cal_init_perturbative_matrix(
        perturbative_matrix_path=perturbative_matrix_path,
        perturbative_matrix_save_path=perturbative_matrix_path,
        **cal_kwargs)

    prtVec = PerturbativeBlockLayer.pert_matrix_to_vec(prtMtrx, maxmn)

    pbl = PerturbativeBlockLayer(2 * maxmn + 1,
                                 init_weight=prtVec,
                                 grouping_triplet=False)  # initial coefficients
    sig = sig[..., 0:symbol_num]
    sig_ori = sig.copy()
    prbs = prbs[..., 0:symbol_num * mod_order]

    sig = torch.from_numpy(sig)
    sig = torch.view_as_real(sig)
    sig = CR_DSPPytorch.norm_power(sig)
    sig = sig[np.newaxis, ...]

    sig_len = sig.shape[-2]
    perturbative_terms = torch.zeros([2, sig_len, 2])

    step = 512
    block_size = step + maxmn + 2
    bs = util.BlockSelector(sample_num=sig_len, block_size=block_size, step=step)

    with torch.no_grad():
        for indx in range(len(bs)):
            sig_block = sig[..., bs[indx], :]
            perturbative_terms[:, indx * step:(indx + 1) * step, :] = pbl(sig_block)[...,
                                                                      bs.pre_overhead:bs.pre_overhead + step, :]

    perturbative_terms = perturbative_terms.data.cpu().numpy()
    perturbative_terms = perturbative_terms[..., 0] + 1j * perturbative_terms[..., 1]
    sig = sig.squeeze()
    sig = sig[..., 0] + 1j * sig[..., 1]
    sig = sig.data.cpu().numpy()

    def est_signal_performance(xi, phi, eta):
        energy_divergence = 1j * 8 / 9 * gm * sig * Leff * power * 3 / 2 * eta
        eqsig = sig + np.exp(1j * phi) * (xi * perturbative_terms - energy_divergence)
        eqsig = eqsig.squeeze()
        ber = util.pr_ber(eqsig, prbs, constellations=util.CONST_16QAM, mod_order=mod_order)[1]
        Q = np.mean(util.ber2q(ber))
        return Q

    optimizer = BayesianOptimization(
        f=est_signal_performance,
        pbounds={'xi': (0, 1e-3), 'phi': (0, 2 * np.pi), 'eta': (10, 50)},
        verbose=1
    )
    utility = UtilityFunction(kind='ei', xi=0.1)
    optimizer.maximize(init_points=20, n_iter=300, acquisition_function=utility)
    xi_opt = optimizer.max['params']['xi']
    phi_opt = optimizer.max['params']['phi']
    eta_opt = optimizer.max['params']['eta']
    print(f'optimal xi: {xi_opt}, optimal phi:{phi_opt} optimal eta:{eta_opt}')

    '''ERP-PB-NLC techniques'''
    energy_diver = 1j * 8 / 9 * gm * sig * Leff * power * 3 / 2 * eta_opt
    eqsig = sig_ori + (xi_opt * perturbative_terms - energy_diver) * np.exp(1j * phi_opt)
    ber_eq = util.pr_ber(eqsig, prbs, constellations=util.CONST_16QAM, mod_order=mod_order)[1]
    Q_eq = np.mean(util.ber2q(ber_eq))
    print(f' With ERP, the Q factor is {Q_eq}')
    return np.mean(ber_eq), Q_eq


def change_file(filepath):
    import pandas as pd
    df = pd.read_csv(os.path.join(filepath, 'Learned_TD_NLC_vary_tap_and_cluster.csv'))
    df.columns = df.columns.str.strip()

    df_sorted = df.sort_values(by=['cluster', 'tap']).reset_index(drop=True)
    df_sorted.to_csv(os.path.join(filepath, 'Learned_TD_NLC_vary_cluster_and_tap.csv'), index=False)
    print(f'result saved at \"{filepath}\"')




if __name__ == '__main__':
    warnings.filterwarnings("ignore", category=UserWarning)
    lp_range = list(np.arange(-2, 4, 1))
    init_method = 'TDS'
    tap_list = [427, 447, 507, 557, 607, 657, 707]
    # ERP_NLC_with_TDC(lp_range=[2], tap_list=tap_list, init_method=init_method,
    #                  save_result_flag=True, clusters_flag=False, quantization_flag=False)
    TDC_Cal_Complexity(lp=2, tap_list=tap_list, cluster_list=[12, 14, 16, 18, 20, 22, 24], init_method=init_method,
                       save_result_flag=True)
    filepath = os.path.join(BASE_DIR, 'result/ERP_NLC_with_TDC')
    change_file(filepath)